import matplotlib.pyplot as plt
import numpy as np
import seaborn as sns
from matplotlib.ticker import FuncFormatter

# 设置学术图表风格
plt.rcParams.update({
    'font.family': 'serif',
    'font.serif': ['Times New Roman'],
    'font.size': 12,
    'axes.labelsize': 14,
    'axes.titlesize': 16,
    'legend.fontsize': 12,
    'xtick.labelsize': 12,
    'ytick.labelsize': 12,
    'figure.figsize': (10, 6),
    'figure.dpi': 300,
    'axes.grid': False,
    'grid.linestyle': '--',
    'grid.alpha': 0.3
})
sns.set_style("whitegrid")

# 数据准备
methods = ['FedAvg', 'FedMe', 'FL+DB', 'MFL-NoPrune', 'MFL-NoCluster', 'MFL (Ours)']
comm_cost = [2450, 1820, 1530, 1020, 980, 620]  # 单位: MB
converge_rounds = [120, 95, 85, 70, 65, 52]

# 创建图表和双坐标轴
fig, ax1 = plt.subplots()
ax2 = ax1.twinx()

# 颜色设置
bar_color = sns.color_palette("Blues_d", len(methods))
line_color = sns.color_palette("Reds_d", len(methods))[3]

# 柱状图 - 通信开销
bars = ax1.bar(methods, comm_cost, color=bar_color, alpha=0.85, edgecolor='black')
ax1.set_ylabel('Communication Cost (MB)')
ax1.set_ylim(0, 2600)

# 折线图 - 收敛轮次
line = ax2.plot(methods, converge_rounds, 
               marker='o', markersize=8, 
               color=line_color, linewidth=3, 
               label='Convergence Rounds')
ax2.set_ylabel('Convergence Rounds')
ax2.set_ylim(0, 130)

# 添加数据标签
for i, (cost, rounds) in enumerate(zip(comm_cost, converge_rounds)):
    ax1.text(i, cost + 50, f'{cost} MB', 
            ha='center', va='bottom', fontsize=10)
    ax2.text(i, rounds + 3, f'{rounds}', 
            ha='center', va='bottom', fontsize=10, color=line_color)

# 添加创新点标注
# ax1.text(5.2, 2000, r'$\downarrow$ 75% Comm Cost vs FedAvg', 
#          fontsize=12, bbox=dict(facecolor='lightyellow', alpha=0.8))
# ax2.text(5.2, 40, r'$\downarrow$ 56.7% Converge Time', 
#          fontsize=12, bbox=dict(facecolor='mistyrose', alpha=0.8))

# 添加技术说明箭头
# ax1.annotate('Dynamic Pruning\nContribution: ↓38%', 
#              xy=(3, 1020), xytext=(1.5, 1800),
#              arrowprops=dict(arrowstyle='->', lw=1.5, color='darkblue'),
#              bbox=dict(boxstyle="round,pad=0.3", fc="lightblue", ec="steelblue", alpha=0.8))

# ax1.annotate('DBSCAN Optimization\nContribution: ↓36.6%', 
#              xy=(5, 620), xytext=(3.5, 1500),
#              arrowprops=dict(arrowstyle='->', lw=1.5, color='darkgreen'),
#              bbox=dict(boxstyle="round,pad=0.3", fc="lightgreen", ec="darkgreen", alpha=0.8))

# 格式设置
ax1.yaxis.set_major_formatter(FuncFormatter(lambda x, _: f'{x:,.0f}'))
# ax1.set_title('Communication Cost vs Convergence Speed Comparison', fontweight='bold', pad=20)
ax1.set_xlabel('Federated Learning Methods')

# 添加图例
from matplotlib.patches import Patch
legend_elements = [
    Patch(facecolor=bar_color[0], edgecolor='black', label='Comm Cost'),
    plt.Line2D([0], [0], marker='o', color='w', markerfacecolor=line_color, markersize=10, label='Converge Rounds')
]
ax1.legend(handles=legend_elements, loc='upper right', framealpha=0.9)

# 学术图表规范调整
plt.tight_layout()
plt.subplots_adjust(top=0.9)
fig.patch.set_facecolor('white')

# 保存为高分辨率图片
plt.savefig('Fig2_Comm_vs_Convergence.png', dpi=300, bbox_inches='tight')
plt.show()
